#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
代码提取器 - 用于软著申请
提取指定目录下的所有代码文件，并将它们整合成一个文档
支持TXT、Markdown和Word格式输出
"""

import os
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import re
import datetime
import shutil
from pathlib import Path

# 尝试导入python-docx库，用于生成Word文档
try:
    import docx
    from docx.shared import Pt, Inches, Cm
    from docx.enum.text import WD_PARAGRAPH_ALIGNMENT, WD_LINE_SPACING
    HAS_DOCX = True
except ImportError:
    HAS_DOCX = False

# 中文五号字体大小（10.5磅）
CN_FONT_SIZE_5 = 10.5

# 默认字体大小（9磅，小五号）
DEFAULT_FONT_SIZE = 9.0

class CodeExtractor:
    def __init__(self, root):
        self.root = root
        self.root.title("代码提取器 - 软著申请工具")
        self.root.geometry("800x600")
        self.root.resizable(True, True)
        
        self.source_dir = tk.StringVar()
        self.output_dir = tk.StringVar()
        self.file_extensions = tk.StringVar(value=".py .java .js .html .css .cpp .c .h .cs .php")
        self.exclude_dirs = tk.StringVar(value="node_modules,venv,__pycache__,.git,.idea")
        self.exclude_files = tk.StringVar(value="")
        self.output_format = tk.StringVar(value="txt")
        self.include_comments = tk.BooleanVar(value=False)
        self.include_file_info = tk.BooleanVar(value=False)
        self.include_header = tk.BooleanVar(value=False)
        self.font_size = tk.DoubleVar(value=DEFAULT_FONT_SIZE)
        self.remove_empty_lines = tk.BooleanVar(value=True)  # 默认删除空行
        
        self._create_widgets()
    
    def _create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 源目录选择
        ttk.Label(main_frame, text="源代码目录:").grid(row=0, column=0, sticky=tk.W, pady=5)
        ttk.Entry(main_frame, textvariable=self.source_dir, width=50).grid(row=0, column=1, sticky=tk.W+tk.E, pady=5)
        ttk.Button(main_frame, text="浏览...", command=self._browse_source_dir).grid(row=0, column=2, padx=5, pady=5)
        
        # 输出目录选择
        ttk.Label(main_frame, text="输出目录:").grid(row=1, column=0, sticky=tk.W, pady=5)
        ttk.Entry(main_frame, textvariable=self.output_dir, width=50).grid(row=1, column=1, sticky=tk.W+tk.E, pady=5)
        ttk.Button(main_frame, text="浏览...", command=self._browse_output_dir).grid(row=1, column=2, padx=5, pady=5)
        
        # 输出格式选择
        ttk.Label(main_frame, text="输出格式:").grid(row=2, column=0, sticky=tk.W, pady=5)
        format_frame = ttk.Frame(main_frame)
        format_frame.grid(row=2, column=1, sticky=tk.W, pady=5)
        
        ttk.Radiobutton(format_frame, text="TXT文本", value="txt", variable=self.output_format).pack(side=tk.LEFT, padx=10)
        ttk.Radiobutton(format_frame, text="Markdown", value="md", variable=self.output_format).pack(side=tk.LEFT, padx=10)
        
        # 如果支持docx，添加Word选项
        if HAS_DOCX:
            ttk.Radiobutton(format_frame, text="Word文档", value="docx", variable=self.output_format).pack(side=tk.LEFT, padx=10)
        else:
            ttk.Label(format_frame, text="(安装python-docx库以启用Word输出)").pack(side=tk.LEFT, padx=10)
        
        # 选项框架
        options_frame = ttk.LabelFrame(main_frame, text="输出选项")
        options_frame.grid(row=3, column=0, columnspan=3, sticky=tk.W+tk.E, pady=5)
        
        # 注释选项
        ttk.Checkbutton(options_frame, text="包含注释", variable=self.include_comments).grid(row=0, column=0, sticky=tk.W, padx=10, pady=5)
        
        # 文件信息选项
        ttk.Checkbutton(options_frame, text="包含文件路径信息", variable=self.include_file_info).grid(row=0, column=1, sticky=tk.W, padx=10, pady=5)
        
        # 报告头部选项
        ttk.Checkbutton(options_frame, text="包含报告头部信息", variable=self.include_header).grid(row=0, column=2, sticky=tk.W, padx=10, pady=5)
        
        # 删除空行选项
        ttk.Checkbutton(options_frame, text="删除空行", variable=self.remove_empty_lines).grid(row=1, column=0, sticky=tk.W, padx=10, pady=5)
        
        # 字体大小选项（仅适用于Word格式）
        font_frame = ttk.Frame(options_frame)
        font_frame.grid(row=1, column=1, columnspan=2, sticky=tk.W, padx=10, pady=5)
        
        ttk.Label(font_frame, text="字体大小(磅):").pack(side=tk.LEFT, padx=5)
        font_sizes = [7.0, 8.0, 9.0, 10.0, 10.5, 12.0]
        font_combo = ttk.Combobox(font_frame, textvariable=self.font_size, values=font_sizes, width=5)
        font_combo.pack(side=tk.LEFT, padx=5)
        ttk.Label(font_frame, text="(仅适用于Word格式)").pack(side=tk.LEFT, padx=5)
        
        # 文件扩展名
        ttk.Label(main_frame, text="文件扩展名:").grid(row=4, column=0, sticky=tk.W, pady=5)
        ttk.Entry(main_frame, textvariable=self.file_extensions, width=50).grid(row=4, column=1, sticky=tk.W+tk.E, columnspan=2, pady=5)
        ttk.Label(main_frame, text="(用空格分隔，例如: .py .java .js)").grid(row=5, column=1, sticky=tk.W, pady=0)
        
        # 排除目录
        ttk.Label(main_frame, text="排除目录:").grid(row=6, column=0, sticky=tk.W, pady=5)
        ttk.Entry(main_frame, textvariable=self.exclude_dirs, width=50).grid(row=6, column=1, sticky=tk.W+tk.E, columnspan=2, pady=5)
        ttk.Label(main_frame, text="(用逗号分隔，例如: node_modules,venv)").grid(row=7, column=1, sticky=tk.W, pady=0)
        
        # 排除文件
        ttk.Label(main_frame, text="排除文件:").grid(row=8, column=0, sticky=tk.W, pady=5)
        ttk.Entry(main_frame, textvariable=self.exclude_files, width=50).grid(row=8, column=1, sticky=tk.W+tk.E, columnspan=2, pady=5)
        ttk.Label(main_frame, text="(用逗号分隔，例如: test.py,temp.js)").grid(row=9, column=1, sticky=tk.W, pady=0)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        ttk.Label(main_frame, text="进度:").grid(row=10, column=0, sticky=tk.W, pady=10)
        self.progress_bar = ttk.Progressbar(main_frame, variable=self.progress_var, maximum=100)
        self.progress_bar.grid(row=10, column=1, sticky=tk.W+tk.E, columnspan=2, pady=10)
        
        # 状态标签
        self.status_var = tk.StringVar(value="就绪")
        self.status_label = ttk.Label(main_frame, textvariable=self.status_var)
        self.status_label.grid(row=11, column=0, columnspan=3, sticky=tk.W, pady=5)
        
        # 日志文本框
        ttk.Label(main_frame, text="日志:").grid(row=12, column=0, sticky=tk.W, pady=5)
        log_frame = ttk.Frame(main_frame)
        log_frame.grid(row=13, column=0, columnspan=3, sticky=tk.W+tk.E+tk.N+tk.S, pady=5)
        
        self.log_text = tk.Text(log_frame, height=10, width=80, wrap=tk.WORD)
        scrollbar = ttk.Scrollbar(log_frame, command=self.log_text.yview)
        self.log_text.configure(yscrollcommand=scrollbar.set)
        
        self.log_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 按钮框架
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=14, column=0, columnspan=3, pady=10)
        
        ttk.Button(button_frame, text="开始提取", command=self.extract_code).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="退出", command=self.root.quit).pack(side=tk.LEFT, padx=5)
        
        # 设置列权重
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(13, weight=1)
    
    def _browse_source_dir(self):
        directory = filedialog.askdirectory(title="选择源代码目录")
        if directory:
            self.source_dir.set(directory)
    
    def _browse_output_dir(self):
        directory = filedialog.askdirectory(title="选择输出目录")
        if directory:
            self.output_dir.set(directory)
    
    def _log(self, message):
        """添加日志消息到日志文本框"""
        self.log_text.insert(tk.END, f"{message}\n")
        self.log_text.see(tk.END)
        self.root.update_idletasks()
    
    def _detect_encoding(self, file_path):
        """检测文件编码，不使用chardet库"""
        encodings = ['utf-8', 'gbk', 'gb2312', 'ascii', 'latin-1', 'big5']
        for encoding in encodings:
            try:
                with open(file_path, 'r', encoding=encoding) as f:
                    f.read()
                return encoding
            except UnicodeDecodeError:
                continue
        return 'utf-8'  # 如果都失败，默认返回utf-8
    
    def _get_output_filename(self, project_name):
        """根据选择的格式生成输出文件名"""
        output_format = self.output_format.get()
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"{project_name}_代码提取_{timestamp}.{output_format}"
    
    def _remove_comments(self, content, file_ext):
        """根据文件类型移除注释"""
        if not self.include_comments.get():
            # 移除C风格注释（C, C++, Java, JavaScript, PHP等）
            if file_ext in ['.c', '.cpp', '.h', '.cs', '.java', '.js', '.php']:
                # 移除多行注释 /* ... */
                content = re.sub(r'/\*[\s\S]*?\*/', '', content)
                # 移除单行注释 //
                content = re.sub(r'//.*', '', content)
            
            # 移除Python注释
            elif file_ext in ['.py']:
                # 移除多行文档字符串 """ ... """ 或 ''' ... '''
                content = re.sub(r'"""[\s\S]*?"""', '', content)
                content = re.sub(r"'''[\s\S]*?'''", '', content)
                # 移除单行注释 #
                content = re.sub(r'#.*', '', content)
            
            # 移除HTML/XML注释
            elif file_ext in ['.html', '.htm', '.xml']:
                content = re.sub(r'<!--[\s\S]*?-->', '', content)
            
            # 移除CSS注释
            elif file_ext in ['.css']:
                content = re.sub(r'/\*[\s\S]*?\*/', '', content)
            
            # 移除连续的空行（两个或更多换行符）
            content = re.sub(r'\n\s*\n', '\n\n', content)
        
        # 如果需要删除所有空行
        if self.remove_empty_lines.get():
            # 移除所有空行（只包含空白字符的行）
            content = re.sub(r'\n\s*\n', '\n', content)
            # 再次处理，确保所有空行都被删除
            content = '\n'.join([line for line in content.split('\n') if line.strip()])
        
        return content
    
    def _write_txt_file(self, all_files, output_file, source_dir):
        """将代码写入TXT文件"""
        with open(output_file, 'w', encoding='utf-8') as out_file:
            # 写入标题和时间戳（如果需要）
            if self.include_header.get():
                out_file.write(f"代码提取报告\n")
                out_file.write(f"生成时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                out_file.write(f"源目录: {source_dir}\n")
                out_file.write("=" * 80 + "\n\n")
            
            for i, file_path in enumerate(all_files):
                rel_path = os.path.relpath(file_path, source_dir)
                self._log(f"处理文件 ({i+1}/{len(all_files)}): {rel_path}")
                
                # 更新进度条
                progress = (i + 1) / len(all_files) * 100
                self.progress_var.set(progress)
                
                # 写入文件分隔符和文件名（如果需要）
                if self.include_file_info.get():
                    out_file.write(f"\n{'=' * 80}\n")
                    out_file.write(f"文件: {rel_path}\n")
                    out_file.write(f"{'-' * 80}\n\n")
                elif i > 0:  # 如果不包含文件信息，但不是第一个文件，添加一个分隔符
                    out_file.write("\n\n")
                
                try:
                    # 检测文件编码
                    encoding = self._detect_encoding(file_path)
                    
                    # 读取文件内容
                    with open(file_path, 'r', encoding=encoding, errors='replace') as in_file:
                        content = in_file.read()
                        
                        # 移除注释
                        file_ext = os.path.splitext(file_path)[1]
                        content = self._remove_comments(content, file_ext)
                        
                        # 写入处理后的内容
                        out_file.write(content)
                except Exception as e:
                    self._log(f"处理文件 {rel_path} 时出错: {str(e)}")
                    if self.include_file_info.get():
                        out_file.write(f"[无法读取文件内容: {str(e)}]\n")
    
    def _write_markdown_file(self, all_files, output_file, source_dir):
        """将代码写入Markdown文件"""
        with open(output_file, 'w', encoding='utf-8') as out_file:
            # 写入标题和时间戳（如果需要）
            if self.include_header.get():
                out_file.write(f"# 代码提取报告\n\n")
                out_file.write(f"**生成时间**: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
                out_file.write(f"**源目录**: {source_dir}\n\n")
                out_file.write("---\n\n")
            
            # 创建目录（如果需要）
            if self.include_file_info.get():
                out_file.write("## 目录\n\n")
                for i, file_path in enumerate(all_files):
                    rel_path = os.path.relpath(file_path, source_dir)
                    file_link = rel_path.replace(" ", "%20")
                    out_file.write(f"{i+1}. [{rel_path}](#{file_link.replace('/', '-').replace('.', '-')})\n")
                out_file.write("\n---\n\n")
            
            for i, file_path in enumerate(all_files):
                rel_path = os.path.relpath(file_path, source_dir)
                self._log(f"处理文件 ({i+1}/{len(all_files)}): {rel_path}")
                
                # 更新进度条
                progress = (i + 1) / len(all_files) * 100
                self.progress_var.set(progress)
                
                # 写入文件标题和链接锚点（如果需要）
                if self.include_file_info.get():
                    anchor = rel_path.replace("/", "-").replace(".", "-")
                    out_file.write(f"## {rel_path} <a id='{anchor}'></a>\n\n")
                elif i > 0:  # 如果不包含文件信息，但不是第一个文件，添加一个分隔符
                    out_file.write("\n\n")
                
                try:
                    # 检测文件编码
                    encoding = self._detect_encoding(file_path)
                    
                    # 读取文件内容
                    with open(file_path, 'r', encoding=encoding, errors='replace') as in_file:
                        content = in_file.read()
                        
                        # 移除注释
                        file_ext = os.path.splitext(file_path)[1]
                        content = self._remove_comments(content, file_ext)
                        
                        # 获取文件扩展名用于代码块语法高亮
                        ext = os.path.splitext(file_path)[1][1:]
                        
                        # 如果包含文件信息，使用代码块格式；否则直接写入代码
                        if self.include_file_info.get():
                            out_file.write(f"```{ext}\n{content}\n```\n\n")
                        else:
                            out_file.write(f"{content}")
                except Exception as e:
                    self._log(f"处理文件 {rel_path} 时出错: {str(e)}")
                    if self.include_file_info.get():
                        out_file.write(f"*无法读取文件内容: {str(e)}*\n\n")
    
    def _write_word_file(self, all_files, output_file, source_dir):
        """将代码写入Word文档，确保每页至少50行代码"""
        if not HAS_DOCX:
            self._log("未安装python-docx库，无法生成Word文档")
            messagebox.showerror("错误", "未安装python-docx库，无法生成Word文档\n请使用pip install python-docx安装")
            return False
        
        doc = docx.Document()
        
        # 使用正常页边距（Word默认）
        
        # 添加标题和时间信息（如果需要）
        if self.include_header.get():
            title = doc.add_heading('代码提取报告', level=0)
            title.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
            
            doc.add_paragraph(f"生成时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            doc.add_paragraph(f"源目录: {source_dir}")
            doc.add_paragraph("").add_run("").add_break()
        
        for i, file_path in enumerate(all_files):
            rel_path = os.path.relpath(file_path, source_dir)
            self._log(f"处理文件 ({i+1}/{len(all_files)}): {rel_path}")
            
            # 更新进度条
            progress = (i + 1) / len(all_files) * 100
            self.progress_var.set(progress)
            
            # 添加文件标题（如果需要）
            if self.include_file_info.get():
                heading = doc.add_heading(f"文件: {rel_path}", level=1)
                # 设置标题段落的间距
                heading.paragraph_format.space_before = Pt(6)
                heading.paragraph_format.space_after = Pt(6)
            elif i > 0:  # 如果不包含文件信息，但不是第一个文件，添加一个分隔符
                para = doc.add_paragraph("")
                para.paragraph_format.space_before = Pt(3)
                para.paragraph_format.space_after = Pt(3)
            
            try:
                # 检测文件编码
                encoding = self._detect_encoding(file_path)
                
                # 读取文件内容
                with open(file_path, 'r', encoding=encoding, errors='replace') as in_file:
                    content = in_file.read()
                    
                    # 移除注释
                    file_ext = os.path.splitext(file_path)[1]
                    content = self._remove_comments(content, file_ext)
                    
                    # 添加代码段落
                    code_para = doc.add_paragraph()
                    code_run = code_para.add_run(content)
                    code_run.font.name = 'Times New Roman'  # 使用Times New Roman字体
                    code_run.font.size = Pt(self.font_size.get())  # 使用用户选择的字体大小
                    
                    # 设置段落格式，减小行间距以增加每页行数
                    code_para.paragraph_format.line_spacing_rule = WD_LINE_SPACING.EXACTLY
                    code_para.paragraph_format.line_spacing = Pt(self.font_size.get() * 1.1)  # 行距设为字体大小的1.1倍
                    code_para.paragraph_format.space_before = Pt(0)
                    code_para.paragraph_format.space_after = Pt(0)
                    
                    # 添加分隔符（如果需要）
                    if self.include_file_info.get():
                        sep_para = doc.add_paragraph("")
                        sep_para.paragraph_format.space_before = Pt(3)
                        sep_para.paragraph_format.space_after = Pt(3)
            except Exception as e:
                self._log(f"处理文件 {rel_path} 时出错: {str(e)}")
                if self.include_file_info.get():
                    doc.add_paragraph(f"[无法读取文件内容: {str(e)}]")
        
        # 保存文档
        doc.save(output_file)
        return True
    
    def extract_code(self):
        """提取代码的主函数"""
        source_dir = self.source_dir.get().strip()
        output_dir = self.output_dir.get().strip()
        
        if not source_dir or not os.path.isdir(source_dir):
            messagebox.showerror("错误", "请选择有效的源代码目录")
            return
        
        if not output_dir or not os.path.isdir(output_dir):
            messagebox.showerror("错误", "请选择有效的输出目录")
            return
        
        # 获取文件扩展名列表
        extensions = [ext.strip() for ext in self.file_extensions.get().split() if ext.strip()]
        if not extensions:
            messagebox.showerror("错误", "请指定至少一个文件扩展名")
            return
        
        # 获取排除目录列表
        exclude_dirs = [dir.strip() for dir in self.exclude_dirs.get().split(',') if dir.strip()]
        
        # 获取排除文件列表
        exclude_files = [file.strip() for file in self.exclude_files.get().split(',') if file.strip()]
        
        try:
            self.status_var.set("正在提取代码...")
            self._log(f"开始提取代码，源目录: {source_dir}")
            self._log(f"文件类型: {', '.join(extensions)}")
            self._log(f"排除目录: {', '.join(exclude_dirs)}")
            self._log(f"排除文件: {', '.join(exclude_files)}")
            self._log(f"包含注释: {'是' if self.include_comments.get() else '否'}")
            self._log(f"包含文件信息: {'是' if self.include_file_info.get() else '否'}")
            self._log(f"包含报告头部: {'是' if self.include_header.get() else '否'}")
            self._log(f"删除空行: {'是' if self.remove_empty_lines.get() else '否'}")
            if self.output_format.get() == "docx":
                self._log(f"字体大小: {self.font_size.get()} 磅")
            
            # 收集所有符合条件的文件
            all_files = []
            for root, dirs, files in os.walk(source_dir):
                # 排除指定目录
                dirs[:] = [d for d in dirs if d not in exclude_dirs]
                
                for file in files:
                    if any(file.endswith(ext) for ext in extensions) and file not in exclude_files:
                        all_files.append(os.path.join(root, file))
            
            total_files = len(all_files)
            self._log(f"找到 {total_files} 个文件")
            
            if total_files == 0:
                messagebox.showinfo("信息", "没有找到符合条件的文件")
                self.status_var.set("就绪")
                return
            
            # 获取项目名称（源目录的最后一个文件夹名）
            project_name = os.path.basename(os.path.normpath(source_dir))
            
            # 根据选择的格式生成输出文件名
            output_filename = self._get_output_filename(project_name)
            output_file = os.path.join(output_dir, output_filename)
            
            output_format = self.output_format.get()
            success = False
            
            if output_format == "txt":
                self._write_txt_file(all_files, output_file, source_dir)
                success = True
            elif output_format == "md":
                self._write_markdown_file(all_files, output_file, source_dir)
                success = True
            elif output_format == "docx":
                success = self._write_word_file(all_files, output_file, source_dir)
            else:
                messagebox.showerror("错误", f"不支持的输出格式: {output_format}")
                self.status_var.set("出错")
                return
            
            if success:
                self.status_var.set("代码提取完成")
                self._log(f"代码提取完成，已保存到: {output_file}")
                messagebox.showinfo("完成", f"代码提取完成，已保存到:\n{output_file}")
            
        except Exception as e:
            self.status_var.set("出错")
            self._log(f"错误: {str(e)}")
            messagebox.showerror("错误", f"代码提取过程中出错:\n{str(e)}")

if __name__ == "__main__":
    root = tk.Tk()
    app = CodeExtractor(root)
    root.mainloop() 